import torch
from torch import nn
import time
import sys
import pickle
import numpy as np
from pytorch_grad_cam import GradCAM, ScoreCAM, GradCAMPlusPlus
from pytorch_grad_cam.utils.image import show_cam_on_image
from torchvision import transforms
# from LoadData import load_cifar10, load_fashion_mnist, load_awa2
from Utils import get_awa2_images
import cv2
sys.path.append("./")

class AlexNet(nn.Module):

    def __init__(self, num_class=10, channels=1):
        super(AlexNet, self).__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(channels, 96, 11, 4),
            nn.BatchNorm2d(96),
            nn.ReLU(),
            nn.MaxPool2d(3, 2),

            nn.Conv2d(96, 256, 5, 1, 2),
            nn.BatchNorm2d(256),
            nn.ReLU(),
            nn.MaxPool2d(3, 2),

            nn.Conv2d(256, 384, 3, 1, 1),
            nn.BatchNorm2d(384),
            nn.ReLU(),

            nn.Conv2d(384, 384, 3, 1, 1),
            nn.BatchNorm2d(384),
            nn.ReLU(),

            nn.Conv2d(384, 256, 3, 1, 1),
            nn.BatchNorm2d(256),
            nn.ReLU(),

            nn.MaxPool2d(3, 2)
        )

        self.fc = nn.Sequential(
            nn.Linear(256 * 5 * 5, 4096),
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(4096, 4096),
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(4096, num_class),
        )

    def forward(self, img):
        feature = self.conv(img)
        # print(feature.shape)
        # exit(0)
        output = self.fc(feature.view(img.shape[0], -1))
        return output

model = AlexNet(num_class=50, channels=3)
model.load_state_dict(torch.load("./Model/AwA2/alex_model.pth"))
print(model)
layer = model.conv[17]

cam = GradCAMPlusPlus(model=model, target_layer=layer, use_cuda=True)


origin_images, input_images, labels = get_awa2_images(128, 224)
grayscale_cam = cam(input_tensor=input_images, target_category=labels, aug_smooth=True)
grayscale_cam = grayscale_cam[100, :]
visualization = show_cam_on_image(origin_images[100, :].transpose((1, 2, 0)), grayscale_cam, use_rgb=False)
print(visualization.shape)
print(type(visualization))
cv2.imwrite("testcam.png", visualization)

cam2 = GradCAM(model=model, target_layer=layer, use_cuda=True)
grayscale_cam2 = cam2(input_tensor=input_images, target_category=labels, aug_smooth=True)
grayscale_cam2 = grayscale_cam2[100, :]
visualization2 = show_cam_on_image(origin_images[100, :].transpose((1, 2, 0)), grayscale_cam2, use_rgb=False)
print(visualization2.shape)
print(type(visualization2))
cv2.imwrite("testcam2.png", visualization2)